import json

import cv2
import easydict
import numpy as np

from cjc.tools.logger import Logger
from cjc.tracker.groups import SexGroups, AgeGroups

"""计数器功能"""

x_span = 20

log = Logger("cls_counter")
def create_full_mask_from_shape(shape, span=20):
    m = np.zeros(shape, dtype=np.uint8)
    m[:, span:-span, :] = 255  # 变红  hw
    return m


def create_full_mask(img, span=20):
    m = np.zeros(img.shape, dtype=np.uint8)
    m[:, span:-span, :] = 255  # 变红  hw
    return m


# self.mask = np.zeros(self.img.shape, dtype=np.uint8)
# # img = np.zeros((200, 200, 3), dtype=np.uint8)
# pts = np.array(self.points, np.int32)  # 数据类型必须为 int32
# pts = pts.reshape((1, -1, 2))
# cv2.fillPoly(self.mask, [pts], color=(255, 255, 255))
# cv2.imwrite(self.img_mask_name, self.mask)

def calculate_position(xywh):
    """根据底部坐标判断是否出现过"""
    return [int((xywh[0] + xywh[2]) / 2), int(xywh[3])]


class PersonCounter:
    def __init__(self, im_background=None):
        self.sex_group = SexGroups()
        self.age_group = AgeGroups()
        self.cb = None
        self.pos_dic = {}
        if im_background is not None:
            self.set_mask(im_background)
            self.mask = create_full_mask(self.mask, x_span)
            self.max_x = self.mask.shape[1] - x_span
        else:
            self.mask = None
        self.counter = 0
        self.min_x = x_span

    def check_sex_and_age_detection_finished(self, _id):
        return self.sex_group.check_detect_finished(_id), self.age_group.check_detect_finished(_id)

    def register_new_video_callback(self, cb):
        self.cb = cb

    def set_mask(self, msk):
        if self.mask is None:
            if isinstance(msk, int) or isinstance(msk, str):
                cap = cv2.VideoCapture(msk)
                i = 0
                while i < 10:
                    ret, msk = cap.read()
                    # print("进入set_mask",ret)
                    i += 1
                self.mask = msk
            else:
                self.mask = msk
                self.max_x = self.mask.shape[1] - x_span

    def draw_mask(self, img):
        """把自己的遮罩显示在最终图像上"""
        return cv2.addWeighted(img, alpha=.9, src2=self.mask, beta=0.1, gamma=1)

    def check_in_mask(self, pos):
        assert self.mask is not None, "遮罩不能是空的"
        if pos[0] < self.min_x:
            print("出圈1")
            return False
        if pos[0] > self.max_x:
            print("出圈2")
            return False
        return True

    def get_counter_status(self):
        """计数器状态包括当前人数，总共人数"""
        ed = easydict.EasyDict()
        ed.total = self.counter
        ed.current = len(self.pos_dic)
        ed.sex_group = self.sex_group.get_count_dict()
        ed.age_group = self.age_group.get_count_dict()
        return ed
        # return json.dumps(ed)

    def get_counter_status_for_webapi(self):
        """计数器状态包括当前人数，总共人数"""
        ed = easydict.EasyDict()
        ed.total = self.counter
        # ed.current = len(self.pos_dic)
        sg = self.sex_group.get_count_dict(1)
        ag = self.age_group.get_count_dict(1)
        # 合并字典参照 https://blog.csdn.net/weixin_56842628/article/details/117996305
        ed.update(sg)
        ed.update(ag)
        return ed

    def run(self, id_, bbox):
        """两个数组，分别计算位置，"""
        pos = calculate_position(bbox)
        if self.check_in_mask(pos):  # 检测进入
            if len(self.pos_dic) == 0 and self.cb is not None:
                self.cb()
            if id_ not in self.pos_dic:
                self.pos_dic[id_] = pos
                self.counter += 1
                log.info(f'曾经有过 :{self.counter}人 ')
                log.info(f'当前场景共 :{len(self.pos_dic)}人 ')

        else:  # 检测退出
            if id_ in self.pos_dic:
                del self.pos_dic[id_]

    def clear(self):
        if len(self.pos_dic) > 0:
            self.pos_dic.clear()
        self.sex_group.clear()
        self.age_group.clear()
